import os
import pymysql
from typing import List, Dict, Any, Optional
from jinja2 import Template
from dataclasses import dataclass
from datetime import datetime
from config import settings


@dataclass
class ColumnInfo:
    name: str
    type: str
    nullable: bool
    primary_key: bool
    comment: str
    max_length: int = None


@dataclass
class TableInfo:
    name: str
    comment: str
    columns: List[ColumnInfo]


class DatabaseReader:
    """数据库读取类"""

    def __init__(self, host=None, port=None, user=None, password=None, database=None):
        self.host = host or settings.DB_HOST
        self.port = port or settings.DB_PORT
        self.user = user or settings.DB_USER
        self.password = password or settings.DB_PASSWORD
        self.database = database or settings.DB_NAME

    def get_connection(self):
        """获取数据库连接"""
        return pymysql.connect(
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            database=self.database,
            charset='utf8mb4'
        )

    def get_table_info(self, table_name: str) -> Optional[TableInfo]:
        """从数据库获取表信息"""
        with self.get_connection() as conn:
            # 获取表注释
            with conn.cursor() as cursor:
                cursor.execute(f"""
                    SELECT TABLE_COMMENT 
                    FROM information_schema.TABLES 
                    WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                """, (self.database, table_name))
                table_comment = cursor.fetchone()[0] or table_name

            # 获取列信息
            with conn.cursor() as cursor:
                cursor.execute(f"""
                    SELECT 
                        COLUMN_NAME,
                        DATA_TYPE,
                        IS_NULLABLE,
                        COLUMN_KEY,
                        COLUMN_COMMENT,
                        CHARACTER_MAXIMUM_LENGTH
                    FROM information_schema.COLUMNS 
                    WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                    ORDER BY ORDINAL_POSITION
                """, (self.database, table_name))
                columns = []
                for row in cursor.fetchall():
                    column_name, data_type, is_nullable, column_key, comment, max_length = row
                    columns.append(ColumnInfo(
                        name=column_name,
                        type=self._convert_db_type(data_type),
                        nullable=is_nullable == 'YES',
                        primary_key=column_key == 'PRI',
                        comment=comment or column_name,
                        max_length=max_length
                    ))

        if not columns:
            return None

        return TableInfo(
            name=table_name,
            comment=table_comment,
            columns=columns
        )

    def _convert_db_type(self, db_type: str) -> str:
        """转换数据库类型到标准类型"""
        type_mapping = {
            'int': 'int',
            'bigint': 'bigint',
            'tinyint': 'int',
            'smallint': 'int',
            'mediumint': 'int',
            'varchar': 'varchar',
            'char': 'char',
            'text': 'text',
            'datetime': 'datetime',
            'timestamp': 'datetime',
            'date': 'date',
            'decimal': 'decimal',
            'float': 'decimal',
            'double': 'decimal',
            'boolean': 'boolean',
            'bool': 'boolean'
        }
        return type_mapping.get(db_type.lower(), 'varchar')


class CodeGenerator:
    def __init__(self, table_info: TableInfo, output_path: str = None):
        self.table_info = table_info
        self.class_name = self._convert_to_class_name(table_info.name)
        self.file_name = self.class_name[:1].lower() + self.class_name[1:]
        self.base_path = output_path or os.path.join(os.getcwd(), 'generated_code')

    def _convert_to_class_name(self, table_name: str) -> str:
        """Convert snake_case table name to CamelCase class name"""
        return ''.join(word.capitalize() for word in table_name.split('_'))

    def _get_python_type(self, db_type: str) -> str:
        """Convert database type to Python type"""
        type_mapping = {
            'int': 'int',
            'bigint': 'int',
            'varchar': 'str',
            'char': 'str',
            'text': 'str',
            'datetime': 'datetime',
            'date': 'date',
            'decimal': 'Decimal',
            'boolean': 'bool'
        }
        return type_mapping.get(db_type.lower(), 'str')

    def _get_tortoise_field_type(self, column: ColumnInfo) -> str:
        """Convert database type to Tortoise ORM field type"""
        type_mapping = {
            'int': 'IntField',
            'bigint': 'BigIntField',
            'varchar': 'CharField',
            'char': 'CharField',
            'text': 'TextField',
            'datetime': 'DatetimeField',
            'date': 'DateField',
            'decimal': 'DecimalField',
            'boolean': 'BooleanField'
        }
        field_type = type_mapping.get(column.type.lower(), 'CharField')

        args = []
        if column.primary_key:
            args.append('pk=True')
        if column.max_length and field_type == 'CharField':
            args.append(f'max_length={column.max_length}')
        if column.nullable:
            args.append('null=True')
        if column.comment:
            args.append(f'description="{column.comment}"')

        args_str = ', '.join(args)
        return f'fields.{field_type}({args_str})'

    def generate_model(self) -> str:
        template = '''from tortoise import fields
from datetime import datetime
from .base import BaseModel, TimestampMixin


class {{ class_name }}(BaseModel, TimestampMixin):
    """{{ table_comment }}"""
{% for column in columns %}
    {{ column.name }} = {{ column.field_type }}
{% endfor %}
    class Meta:
        table = "{{ table_name }}"

    def __str__(self):
        return f"{self.{{ primary_key_name }}}"
'''
        columns_info = []
        primary_key_name = None
        for col in self.table_info.columns:
            if col.primary_key:
                primary_key_name = col.name
            columns_info.append({
                'name': col.name,
                'field_type': self._get_tortoise_field_type(col)
            })

        return Template(template).render(
            class_name=self.class_name,
            file_name=self.file_name,
            table_name=self.table_info.name,
            table_comment=self.table_info.comment,
            columns=columns_info,
            primary_key_name=primary_key_name
        )

    def generate_schema(self) -> str:
        template = '''from typing import Optional, List
from pydantic import BaseModel, Field
from datetime import datetime


class BaseSchema(BaseModel):
    class Config:
        allow_population_by_field_name = True
        alias_generator = lambda x: ''.join(word.capitalize() if i else word for i, word in enumerate(x.split('_')))


class {{ class_name }}Base(BaseSchema):
{% for field in fields %}
    {{ field.name }}: Optional[{{ field.type }}] = Field({{ field.default }}, alias="{{ field.alias }}")
{% endfor %}


class {{ class_name }}Create({{ class_name }}Base):
    pass


class {{ class_name }}Update({{ class_name }}Base):
    pass


class {{ class_name }}Delete(BaseModel):
    ids: List[int] = Field(..., description="{{ table_comment }}ID列表")


class PageParams(BaseSchema):
    page_num: int = Field(1, ge=1, alias="pageNum", description="页码")
    page_size: int = Field(10, ge=1, le=100, alias="pageSize", description="每页数量")
{% for field in search_fields %}
    {{ field.name }}: Optional[{{ field.type }}] = Field(None, alias="{{ field.alias }}")
{% endfor %}
'''
        fields = []
        search_fields = []
        for col in self.table_info.columns:
            if not col.primary_key:
                field_info = {
                    'name': col.name,
                    'type': self._get_python_type(col.type),
                    'default': 'None' if col.nullable else '""',
                    'alias': ''.join(word.capitalize() if i else word for i, word in enumerate(col.name.split('_')))
                }
                fields.append(field_info)

                # Add searchable fields
                if col.type.lower() in ['varchar', 'char', 'text']:
                    search_fields.append(field_info)

        return Template(template).render(
            class_name=self.class_name,
            table_comment=self.table_info.comment,
            fields=fields,
            search_fields=search_fields
        )

    def generate_excel(self) -> str:
        template = '''from utils.excel_util import ExcelUtil
from model.{{ model_name }} import {{ class_name }}
from schema.{{ schema_name }} import {{ class_name }}Create
from typing import List, Any, Dict, ClassVar
from io import BytesIO
from datetime import datetime


class {{ class_name }}Excel:
    # Define export columns
    COLUMNS: ClassVar[List[Dict[str, Any]]] = [
{% for column in columns %}
        {"field": "{{ column.field }}", "title": "{{ column.title }}", "required": {{ column.required }}},
{% endfor %}
    ]

    @classmethod
    def get_export_columns(cls) -> List[Dict[str, Any]]:
        return [
{% for column in export_columns %}
            {"field": "{{ column.field }}", "title": "{{ column.title }}"},
{% endfor %}
        ]

    @classmethod
    def get_import_template(cls) -> BytesIO:
        """Get import template"""
        example_data = [{
{% for field in example_data %}
            "{{ field.name }}": {{ field.value }},
{% endfor %}
        }]
        return ExcelUtil.create_import_template(
            columns=cls.COLUMNS,
            example_data=example_data,
            sheet_name="{{ table_comment }}导入模板"
        )

    @classmethod
    async def export_data(cls, data: List[Any]) -> BytesIO:
        """Export data"""
        return await ExcelUtil.export_data(
            data=data,
            columns=cls.get_export_columns(),
            sheet_name="{{ table_comment }}数据"
        )

    @classmethod
    def parse_import_data(cls, file: BytesIO) -> List[Dict[str, Any]]:
        """Parse import data"""
        return ExcelUtil.parse_import_data(
            file=file,
            columns=cls.COLUMNS,
            schema_cls={{ class_name }}Create
        )
'''
        # Generate column definitions
        columns = []
        export_columns = []
        example_data = []

        for col in self.table_info.columns:
            if not col.primary_key:
                # Import columns
                columns.append({
                    'field': col.name,
                    'title': col.comment,
                    'required': str(not col.nullable)
                })

                # Export columns
                export_columns.append({
                    'field': col.name,
                    'title': col.comment
                })

                # Example data
                example_value = self._get_example_value(col)
                if example_value is not None:
                    example_data.append({
                        'name': col.name,
                        'value': example_value
                    })

        return Template(template).render(
            model_name=f"{self.file_name}Model",
            schema_name=f"{self.file_name}Schema",
            class_name=self.class_name,
            table_comment=self.table_info.comment,
            columns=columns,
            export_columns=export_columns,
            example_data=example_data
        )

    def _get_example_value(self, column: ColumnInfo) -> str:
        """Get example value for template based on column type"""
        type_examples = {
            'int': '1',
            'bigint': '1',
            'varchar': '"example"',
            'char': '"A"',
            'text': '"example text"',
            'datetime': 'datetime.now()',
            'date': 'datetime.now().date()',
            'decimal': '1.0',
            'boolean': 'True'
        }
        return type_examples.get(column.type.lower())

    def generate_service(self) -> str:
        template = '''from datetime import datetime
from typing import List, Tuple
from fastapi import HTTPException, UploadFile
from tortoise.contrib.pydantic import pydantic_model_creator
from model.{{ model_name }} import {{ class_name }}
from schema.{{ schema_name }} import {{ class_name }}Create, {{ class_name }}Update, PageParams
from io import BytesIO
from excel.{{ excel_name }} import {{ class_name }}Excel


class PydanticConfig:
    alias_generator = lambda x: ''.join(word.capitalize() if i else word for i, word in enumerate(x.split('_')))
    allow_population_by_field_name = True


{{ class_name }}_Pydantic = pydantic_model_creator({{ class_name }}, name="{{ class_name }}", config_class=PydanticConfig)


class {{ class_name }}Service:
    @staticmethod
    async def create(data: {{ class_name }}Create, creator: str) -> {{ class_name }}:
        """Create new record"""
        data_dict = data.dict(exclude_unset=True)
        data_dict["create_by"] = creator
        data_dict["create_time"] = datetime.now()
        data_dict["status"] = '0'

        obj = await {{ class_name }}.create(**data_dict)
        return await {{ class_name }}_Pydantic.from_tortoise_orm(obj)

    @staticmethod
    async def get_by_id(id: int) -> {{ class_name }}:
        """Get record by ID"""
        obj = await {{ class_name }}.get_or_none({{ primary_key }}=id, status='0')
        if not obj:
            raise HTTPException(status_code=404, detail="{{ table_comment }}不存在")
        return await {{ class_name }}_Pydantic.from_tortoise_orm(obj)

    @staticmethod
    async def update(id: int, data: {{ class_name }}Update, updater: str) -> {{ class_name }}:
        """Update record"""
        obj = await {{ class_name }}.get_or_none({{ primary_key }}=id, status='0')
        if not obj:
            raise HTTPException(status_code=404, detail="{{ table_comment }}不存在")

        update_dict = data.dict(exclude_unset=True)
        update_dict["update_by"] = updater
        update_dict["update_time"] = datetime.now()

        await obj.update_from_dict(update_dict).save()
        return await {{ class_name }}_Pydantic.from_tortoise_orm(obj)

    @staticmethod
    async def delete(ids: List[int], updater: str) -> int:
        """Batch delete records"""
        update_data = {
            "status": '2',
            "update_by": updater,
            "update_time": datetime.now()
        }
        count = await {{ class_name }}.filter({{ primary_key }}__in=ids, status='0').update(**update_data)
        return count

    @staticmethod
    async def get_page(params: PageParams) -> Tuple[List[{{ class_name }}], int]:
        """Get page list"""
        query = {{ class_name }}.filter(status='0')

        # Build query conditions
{% for field in search_fields %}
        if params.{{ field.name }}:
            query = query.filter({{ field.name }}__icontains=params.{{ field.name }})
{% endfor %}
        total = await query.count()

        records = await query.offset((params.page_num - 1) * params.page_size) \
            .limit(params.page_size) \
            .order_by("-create_time")

        return [await {{ class_name }}_Pydantic.from_tortoise_orm(record) for record in records], total

    @staticmethod
    async def export_excel() -> BytesIO:
        """Export all data to Excel"""
        records = await {{ class_name }}.filter(status='0').all()
        return await {{ class_name }}Excel.export_data(records)

    @staticmethod
    def get_import_template() -> BytesIO:
        """Get import template"""
        return {{ class_name }}Excel.get_import_template()

    @staticmethod
    async def import_data(file: UploadFile, creator: str) -> Tuple[int, List[str]]:
        """Import data"""
        content = await file.read()
        file_bytes = BytesIO(content)

        try:
            data_list = {{ class_name }}Excel.parse_import_data(file_bytes)
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))
        except Exception as e:
            raise HTTPException(status_code=400, detail="Excel文件格式错误")

        success_count = 0
        error_msgs = []

        for index, item_data in enumerate(data_list, start=1):
            try:
                item_data['create_by'] = creator
                item_data['create_time'] = datetime.now()
                item_data['status'] = '0'

                await {{ class_name }}.create(**item_data)
                success_count += 1

            except Exception as e:
                error_msgs.append(f"第{index}行: {str(e)}")

        return success_count, error_msgs
'''
        # Get primary key name
        primary_key = next((col.name for col in self.table_info.columns if col.primary_key), None)

        # Get searchable fields
        search_fields = []
        for col in self.table_info.columns:
            if col.type.lower() in ['varchar', 'char', 'text']:
                search_fields.append({'name': col.name})

        return Template(template).render(
            model_name=f"{self.file_name}Model",
            schema_name=f"{self.file_name}Schema",
            excel_name=f"{self.file_name}Excel",
            class_name=self.class_name,
            table_comment=self.table_info.comment,
            primary_key=primary_key,
            search_fields=search_fields
        )

    def generate_controller(self) -> str:
        template = '''from fastapi import APIRouter, Depends, Path, HTTPException, UploadFile, File, Request
from typing import List
from service.{{ service_name }} import {{ class_name }}Service
from schema.{{ schema_name }} import {{ class_name }}Create, {{ class_name }}Update, PageParams
from utils.response import ResponseUtil, ResponseModel
from urllib.parse import quote
from fastapi.responses import StreamingResponse
from datetime import datetime
from io import BytesIO

router = APIRouter(prefix="/{{ router_prefix }}", tags=["{{ table_comment }}"])
service = {{ class_name }}Service()


@router.post(
    "",
    response_model=ResponseModel,
    summary="新增{{ table_comment }}",
    description="创建一个新的{{ table_comment }}。",
    response_description="返回创建成功的信息"
)
async def create(data: {{ class_name }}Create, request: Request):
    """新增{{ table_comment }}"""
    current_user = request.state.user
    creator = current_user["username"]
    result = await service.create(data, creator)
    return ResponseUtil.success(data=result, msg="创建成功")


@router.get(
    "/get/{id}",
    response_model=ResponseModel,
    summary="获取{{ table_comment }}详情",
    description="根据ID获取{{ table_comment }}的详细信息",
    response_description="返回{{ table_comment }}的详细信息"
)
async def get_detail(id: int):
    """获取{{ table_comment }}详情"""
    result = await service.get_by_id(id)
    return ResponseUtil.success(data=result)


@router.patch(
    "/update/{id}",
    response_model=ResponseModel,
    summary="更新{{ table_comment }}",
    description="更新指定ID的{{ table_comment }}信息。支持部分字段更新。",
    response_description="返回更新后的信息"
)
async def update(id: int, data: {{ class_name }}Update, request: Request):
    """更新{{ table_comment }}"""
    current_user = request.state.user
    updater = current_user["username"]
    result = await service.update(id, data, updater)
    return ResponseUtil.success(data=result, msg="更新成功")


@router.delete(
    "/batch/{ids}",
    response_model=ResponseModel,
    summary="批量删除{{ table_comment }}",
    description="根据ID删除{{ table_comment }}，支持多个ID(用逗号分隔)。此操作不可恢复，请谨慎操作。",
    response_description="返回成功删除的数量"
)
async def delete(request: Request, ids: str = Path(..., description="ID，多个用逗号分隔，如：1,2,3")):
    """批量删除{{ table_comment }}"""
    try:
        id_list = [int(id.strip()) for id in ids.split(",")]
    except ValueError:
        raise ResponseUtil.error(code=400, msg="无效的ID格式")

    current_user = request.state.user
    updater = current_user["username"]
    count = await service.delete(id_list, updater)
    return ResponseUtil.success(data={"count": count}, msg=f"成功删除{count}条数据")


@router.get(
    "/page",
    response_model=ResponseModel,
    summary="分页查询{{ table_comment }}列表",
    description="分页获取{{ table_comment }}列表，支持多条件查询",
    response_description="返回分页数据，包含列表数据、总数、页码、每页大小"
)
async def get_page(
    params: PageParams = Depends()
):
    """分页获取{{ table_comment }}列表"""
    records, total = await service.get_page(params)
    return ResponseUtil.success(data={
        "records": records,
        "total": total,
        "page": params.page_num,
        "limit": params.page_size
    })


@router.get(
    "/export",
    summary="导出{{ table_comment }}数据",
    description="导出所有{{ table_comment }}数据到Excel文件",
    response_description="返回Excel文件"
)
async def export_data():
    """导出数据到Excel"""
    excel_data = await service.export_excel()

    filename = f"{{ table_comment }}数据_{datetime.now().strftime('%Y%m%d%H%M%S')}.xlsx"
    encoded_filename = quote(filename)

    headers = {
        'Content-Disposition': f'attachment; filename="{encoded_filename}"; filename*=utf-8\\\'\\\'{encoded_filename}',
        'Content-Type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
    }

    return StreamingResponse(
        BytesIO(excel_data.getvalue()),
        headers=headers,
        media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
    )


@router.get(
    "/importTemplate",
    summary="下载{{ table_comment }}导入模板",
    description="下载{{ table_comment }}导入的Excel模板文件",
    response_description="返回Excel模板文件"
)
async def get_import_template():
    """获取导入模板"""
    template_data = service.get_import_template()
    filename = f"{{ table_comment }}导入模板_{datetime.now().strftime('%Y%m%d%H%M%S')}.xlsx"
    encoded_filename = quote(filename)

    headers = {
        'Content-Disposition': f'attachment; filename="{encoded_filename}"; filename*=utf-8\\\'\\\'{encoded_filename}',
        'Content-Type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
    }

    return StreamingResponse(
        BytesIO(template_data.getvalue()),
        headers=headers,
        media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
    )


@router.post(
    "/import",
    response_model=ResponseModel,
    summary="导入{{ table_comment }}数据",
    description="通过Excel文件批量导入{{ table_comment }}数据",
    response_description="返回导入结果，包括成功数量和错误信息"
)
async def import_data(
    file: UploadFile = File(..., description="Excel文件")
):
    """导入数据"""
    if not file.filename.endswith(('.xls', '.xlsx')):
        return ResponseUtil.error(msg="只支持Excel文件格式（.xls, .xlsx）")

    creator = "admin"

    try:
        success_count, error_msgs = await service.import_data(file, creator)
        return ResponseUtil.success(data={
            "successCount": success_count,
            "errorMsgs": error_msgs
        }, msg=f"成功导入{success_count}条数据")
    except HTTPException as e:
        return ResponseUtil.error(msg=str(e.detail))
    except Exception as e:
        return ResponseUtil.error(msg=f"导入失败: {str(e)}")

'''
        return Template(template).render(
            service_name=f"{self.file_name}Service",
            schema_name=f"{self.file_name}Schema",
            class_name=self.class_name,
            # router_prefix=self.class_name.replace('_', '/'),
            router_prefix=self.file_name,
            table_comment=self.table_info.comment
        )

    def generate_all(self) -> Dict[str, str]:
        """Generate all code files"""
        result = {
            f"model/{self.file_name}Model.py": self.generate_model(),
            f"schema/{self.file_name}Schema.py": self.generate_schema(),
            f"excel/{self.file_name}Excel.py": self.generate_excel(),
            f"service/{self.file_name}Service.py": self.generate_service(),
            f"controller/{self.file_name}Controller.py": self.generate_controller()
        }
        return result

    def save_to_files(self, base_path: str = None):
        """Save generated code to files"""
        if base_path:
            self.base_path = base_path

        generated_files = self.generate_all()
        for file_path, content in generated_files.items():
            # Create directory if it doesn't exist
            full_path = os.path.join(self.base_path, file_path)
            os.makedirs(os.path.dirname(full_path), exist_ok=True)

            # Write file content
            with open(full_path, 'w', encoding='utf-8') as f:
                f.write(content)


# Example usage:
if __name__ == '__main__':
    # 使用方式1：直接从数据库读取表信息
    db_reader = DatabaseReader()
    table_name = "sys_role_dept"  # 替换为实际的表名
    table_info = db_reader.get_table_info(table_name)

    if table_info:
        # 指定输出路径（可选）
        output_path = "C:\\Users\\DELL\\Desktop\\code"  # 替换为实际的输出路径
        generator = CodeGenerator(table_info, output_path)
        generator.save_to_files()
        print(f"Table {table_name} generate success")
    else:
        print(f"Table {table_name} not found in database")

    # 使用方式2：手动创建表信息
    """
    table_info = TableInfo(
        name="example_table",
        comment="示例表",
        columns=[
            ColumnInfo(
                name="id",
                type="int",
                nullable=False,
                primary_key=True,
                comment="主键ID"
            ),
            # 添加更多列...
        ]
    )
    generator = CodeGenerator(table_info, "path/to/output")
    generator.save_to_files()
    """